#/usr/bin/python27
#-*- coding: utf-8 -*-
#  Copyright 2011-2013 wuweihui@cn.ibm.com
#
#  Licensed under the Apache License, Version 2.0 (the "License");
#  you may not use this file except in compliance with the License.
#  You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
#  Unless required by applicable law or agreed to in writing, software
#  distributed under the License is distributed on an "AS IS" BASIS,
#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#  See the License for the specific language governing permissions and
#  limitations under the License.

import sys
sys.setrecursionlimit(3000)
from socket import *
import threading
import time
from CIMConnection import Client
import subprocess
from config import *
import StringIO

cllock = threading.RLock()

class syncClient(object):
   def __init__(self, port=50007):
      """
      1. Listen messages from server
      2. Send hart beat info to server
      3. Send command status to server
      """
      host = ''
      self.port = port
      self.clientlist = []
      self.ck = socket(AF_INET, SOCK_STREAM)
      self.ck.bind((host, port))
      self.ck.listen(1)
      self.cn = None
      self.runflag = True
      
   def Start(self):
      print 'Server started at', self.port
      self.heartthread = threading.Thread(target=self.HeartBeat, args=())
      self.heartthread.start()
      conn, addr = self.ck.accept()
      while True:         
         with cllock:
            print addr, '----- Connected!'
            self.clientlist.append((conn, addr))
            threading.Thread(target=self.WaitCommand, args=(conn,addr)).start()
         conn, addr = self.ck.accept()
         
   def HeartBeat(self):
      while True:
         for conn, addr in self.clientlist:
            try:
               conn.send('Heart Beat')
            except error, e:
               with cllock:
                  print addr, '----- DisConnected!'
                  self.clientlist.remove((conn, addr))
         time.sleep(5)
   
   def WaitCommand(self, conn, addr='None'):
      while True:
         try:
            data = conn.recv(1024)
         except error, e:
            break
         if not data: continue
         with cllock:
            print addr, '-----', data
         conn.send("Received " + data)
         threading.Thread(target=self.ProcessCommand, args=(conn,data)).start()
         #self.ProcessCommand(conn, data)
   
   def ProcessCommand(self, conn, cmd):
      c = cmd.split()[0]
      self.customizedcmds = ['probe', 'probeex', 'stop', 'sync']
      if c not in self.customizedcmds:
         #with cllock:
         #   print 'Unknown Command:', cmd
         e = ''
         out = ''
         try:
            p = subprocess.Popen(cmd, shell=True, bufsize=1024*1024, stdout=subprocess.PIPE)   
            state = p.wait()
            out = p.stdout.read()         
            if state != 0:
               raise RuntimeError("System Command Failed!")
         except Exception, e:
            with cllock:
               print e
         try:
            conn.send(cmd + str(e) + out)
         except error, e:
            pass
         return
      if c == 'probe':
         try:
            cn = Client.CIMClient()
            self.cn = cn
            cn.connect_to_cimagent()
            cn.probe()
            cn.close_connection()
         except Exception, e:
            with cllock:
               print e
            try:
               conn.send("Result Failure " + cmd + '\n' + str(e))
            except error, e:
               pass
            return
      if c == 'probeex':
         if len(cmd.split()) > 1 and cmd.split()[1].isdigit():
            tot = int(cmd.split()[1])
         else:
            tot = 50
         count = 0
         self.runflag = True
         while self.runflag and count < tot:
            try:
               cn = Client.CIMClient()
               self.cn = cn
               cn.connect_to_cimagent()
               cn.probe()
               cn.close_connection()
            except Exception, e:
               with cllock:
                  print e
               try:
                  conn.send("Result Failure " + cmd + '\n' + str(e))
               except error, e:
                  pass
               count += 1
               cn.close_connection()
               continue
            count += 1
      if c == 'stop':
#         try:
#            conn.send("Stopping...")
#         except error, e:
#            pass
         if self.cn:
            self.cn.close_connection()
            self.cn=None
         self.runflag = False
      if c == 'sync':
         data = cmd[len('sync '):]
         config.save_config(data)
      try:
         conn.send("Result Success " + cmd)
      except error, e:
         pass

class syncServer(object):
   def __init__(self):
      """
      1. connect to client
      2. check health state
      3. send commands
      """
      self.ck = socket(AF_INET, SOCK_STREAM)
      
   def Connect(self, host, port=50007, healthhandle=None):
      self.host = host
      self.ck.connect((host, port))
      self.health = 2
      self.status = 'idle'
      self.healththread = threading.Thread(target=self.CheckHealth, args=(healthhandle,))
      self.healththread.start()
      self.recvthread = threading.Thread(target=self.Recv, args=())
      self.recvthread.start()
      
      if host not in ['localhost', '127.0.0.1']:
         buf = StringIO.StringIO()
         config.write(buf)
         self.SendCmd('sync ' + buf.getvalue())
      
   def DisConnect(self):
      self.ck.close()
      self.health = False
      
   def Recv(self):
      while True:
         data = self.ck.recv(1024)
         if not data: continue         
         if data == 'Heart Beat':
            self.health = 2
         elif data.startswith('Received '):
            with cllock:
               print "Command Deployed!"
         elif data.startswith('Result '):
            with cllock:
               print "Command Finished!"
            self.status = 'idle ' + data[len('Result '):]
            continue
         else:
            with cllock:
               print data
            self.status = 'idle'
      
   def CheckHealth(self, handle=None):
      while True:
         if self.health == 2:
            self.health = 1
         else:
            self.ck.close()
            self.health = 0
            if handle:
               handle(self)
            break
         time.sleep(10) 
      
   def SendCmd(self, cmd):
      self.ck.send(cmd)
      self.status = 'running'
   
   def Shell(self):
      print "Sync server shell started, please input your command!"
      while True:
         cmd = raw_input('')
         if cmd.lower() == "exit":
            break
         self.SendCmd(cmd)
         
class syncServerController(object):
   """
   Need to do:
   1. track status of each session
   2. support remote interruption
   3. can be operate on one session at a time
   """
   def __init__(self):
      self.hostlist = []
      self.Shell()
   
   def Shell(self):
      helpinfo = """Sync server shell started, please input your command!
      help     --  print help info
      add      --  <server ip>  add a server into manage
      remove   --  <server ip>  remove a server
      send     --  <cmd> [<arg1> <arg2> ...]  send a command to run remotely  
      list     --  print out the status of each server
      sync     --  syncronize configuration
      exit
      """
      print helpinfo
      while True:
         cmd = raw_input('')
         if cmd.lower() == "exit":
            break
         elif cmd.lower() == "help":
            print helpinfo
            continue
         elif cmd.lower().startswith('add'):
            try:
               host = cmd.split()[1]
            except:
               print helpinfo
               continue
            try:
               s = syncServer()
               s.Connect(host, healthhandle=self.hostlist.remove)
            except:
               print 'Unable connect to host'
               continue
            self.hostlist.append(s)
         elif cmd.lower().startswith('remove'):
            try:
               host = cmd.split()[1]
            except:
               print helpinfo
               continue
            for s in self.hostlist:
               if s.host == host:
                  self.hostlist.remove(s)
                  try:
                     s.DisConnect()
                  except error, e:
                     pass
                  break
         elif cmd.lower().startswith('send'):
            try:
               cmd = cmd[len('send '):]
            except:
               print helpinfo
               continue
            for s in self.hostlist:
               try:
                  s.SendCmd(cmd)
               except error:
                  print s.host, "Is Dead!"
                  self.hostlist.remove(s)
         elif cmd.lower().startswith('list'):
            for s in self.hostlist:
               print s.host, s.status
         elif cmd.lower().startswith('sync'):
            for s in self.hostlist:
               if s.host not in ['localhost', '127.0.0.1']:
                  config.reload_config()
                  buf = StringIO.StringIO()
                  config.write(buf)
                  for s in self.hostlist:
                     try:
                        s.SendCmd('sync ' + buf.getvalue())
                     except error:
                        print s.host, "Is Dead!"
                        self.hostlist.remove(s)
         else:
            print helpinfo

if __name__ == "__main__":   
   print """
   Copyright 2011-2013 wuweihui@cn.ibm.com
   """
   if len(sys.argv) > 1:
      host = sys.argv[1]
      if host == "ctrl":
         s = syncServerController()
      else:
         s = syncServer()
         s.Connect(host)
         s.Shell()
   else:
      c = syncClient()
      c.Start()